#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <openssl/sha.h>
#include "ovmf_hash_table.h"
#include "utils.h"

/**
 * @brief sha256 function
 *
 * @param [in]  buf   Input data memory address
 * @param [in]  len   Input data length
 * @param [out] len   sha256 result
 *
 * Return: 1 when equal otherwise 0
 */
static int sha256_hash(uint8_t* buf, size_t len, uint8_t *result)
{
    SHA256_CTX ctx;
    // Calculate the hash
    if (SHA256_Init(&ctx) != 1)
        return -1;

    if ((len > 0) && (buf != NULL) && (SHA256_Update(&ctx, buf, len) != 1))
        return -1;

    if (SHA256_Final((uint8_t*)result, &ctx) != 1)
        return -1;

    return 0;
}

/**
 * @brief check if the 2 uuids equal.
 *
 * @param [in] lhv first uuid
 * @param [in] rhv second uuid
 *
 * Return: 1 when equal otherwise 0
 */
static int qemu_uuid_is_equal(const uuid *lhv, const uuid *rhv)
{
    return memcmp(lhv, rhv, sizeof(uuid)) == 0;
}

/**
 * @brief check if the uuid string is valid.
 *
 * The valid regular expression is
 * "[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}"
 *
 * @param [in] str The uuid string to check
 *
 * Return: true if valid otherwise false
 */
static bool qemu_uuid_is_valid(const char *str)
{
    int i;

    for (i = 0; i < (int)strlen(str); i++) {
        const char c = str[i];
        if (i == 8 || i == 13 || i == 18 || i == 23) {
            if (str[i] != '-') {
                return false;
            }
        } else {
            if ((c >= '0' && c <= '9') ||
                (c >= 'A' && c <= 'F') ||
                (c >= 'a' && c <= 'f')) {
                continue;
            }
            return false;
        }
    }
    return i == 36;
}

/**
 * @brief check if the uuid string is valid.
 *
 * The valid regular expression is
 * "[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}"
 *
 * @param [in] str The uuid string to check
 *
 * Return: true if valid otherwise false
 */
static int qemu_uuid_parse(const char *str, uuid *uuid)
{
    unsigned char *uu = &uuid->data[0];
    int ret;

    if (!qemu_uuid_is_valid(str)) {
        return -1;
    }

    ret = sscanf(str, UUID_FMT, &uu[0], &uu[1], &uu[2], &uu[3],
                 &uu[4], &uu[5], &uu[6], &uu[7], &uu[8], &uu[9],
                 &uu[10], &uu[11], &uu[12], &uu[13], &uu[14],
                 &uu[15]);

    if (ret != 16) {
        return -1;
    }
    return 0;
}

/**
 * @brief Swap from uuid format endian (Big Endian) to the opposite or vice versa.
 *
 * @param [in] id The uuid to be swapped
 *
 * Return: The swapped id value
 */
static uuid qemu_uuid_bswap(uuid id)
{
    reverse_bytes((uint8_t *)&id.fields.time_low, sizeof(uint32_t));
    reverse_bytes((uint8_t *)&id.fields.time_mid, sizeof(uint16_t));
    reverse_bytes((uint8_t *)&id.fields.time_high_and_version,
                       sizeof(uint16_t));
    return id;
}

/**
 * @brief Find the data associated with an entry in OVMF's GUIDed table.
 *
 * @param [in] ovmf ovmf The pointer to memory corresponding to the OVMF flash file
 * @param [in] ovmf_size The size of the OVMF flash file
 * @param [in] target_guid_str GUID string of the target
 * @param [in/out] data Filled with a pointer to the entry's value (if not NULL)
 * @param [in/out] data_len Filled with the length of the entry's value (if not NULL)
 *
 * Return: true if the target GUILD entry is found; false otherwise.
 */
/*
 *  The OVMD GUID table structure overview (view from high address to low address)
 *  The GUID table entry data structure is
 *   arbitrary length data
 *   2 byte length of entire entry (count in length itself and GUID)
 *   16 bytes GUID
 *
 *  Assure there are N entries in the GUID table, the format is like:
 *
 *  |     End of OVMF (high memory address)
 *  |     fixed 32 bytes                    |
 *  |     FOOTER_GUID (16 bytes)            | <--- end of GUID table
 *  |     GUID Table Total Length (2 bytes) |
 *  |     Entry N-1 ID (16 bytes)           |
 *  |     Entry N-1 length (2 bytes)        |
 *  |     Entry N-1 Data                    |
 *  |     Entry N-2 ID (16 bytes)           |
 *  |     Entry N-2 length (2 bytes)        |
 *  |     Entry N-2 Data                    |
 *  |     ... ...                           |
 *  |     Entry 0 ID (16 bytes)             |
 *  |     Entry 0 length (2 bytes)          |
 *  |     Entry 0 Data                      | <--  beginning of GUID table
 *  |     (low memory address)              |
 *
 */

#define  FIX_LENGTH_TO_END_OF_OVMF    32
#define  PAGE_SIZE                    4096

bool ovmf_table_find(uint8_t *ovmf, size_t ovmf_size,
                     char *target_guid_str,
                     uint8_t **data, int *data_len)
{
    uint8_t *ptr = NULL;
    int tot_len  = 0;
    uuid footer_guid = {0};
    uuid target_guid = {0};
    int footer_guild_len = 0;

    if (!ovmf || ovmf_size < PAGE_SIZE) {
        return false;
    }

    qemu_uuid_parse(OVMF_TABLE_FOOTER_GUID, &footer_guid);
    footer_guid = qemu_uuid_bswap(footer_guid); /* guids are LE */
    ptr = ovmf + ovmf_size - FIX_LENGTH_TO_END_OF_OVMF;
    ptr -= sizeof(uuid);

    /* check whether @ptr points to FOOTER_GUID */
    if (!qemu_uuid_is_equal((uuid *)ptr, &footer_guid)) {
        return false;
    }

    /* Have found FOOTER_GUID, @ptr points to GUID Table Total Length */
    ptr -= sizeof(uint16_t);

     /*
     * The effective total length should redcue the OVMF_TABLE_FOOTER_GUID entry length
     */
    footer_guild_len = sizeof(uuid) + sizeof(uint16_t);
    tot_len = *(uint16_t *)ptr - footer_guild_len;

    if (tot_len < 0 || tot_len > (ptr - ovmf)) {
        printf("Error: OVMF table has invalid size %d\n", tot_len);
        return false;
    }

    if (tot_len == 0) {
        /* no entries in the OVMF table */
        printf("Error: OVMF table has no entry \n");
        return false;
    }

    /*
     * By now, we have found footer table,
     * @ptr points to the top of OVMF table (end of entry N-1)
     * [@ptr - @tot_len, @ptr) is the region of OVMF table.
     */

    if (qemu_uuid_parse(target_guid_str, &target_guid) < 0) {
        return false;
    }

    target_guid = qemu_uuid_bswap(target_guid); /* guids are LE */
    while (tot_len >= (int)(sizeof(uuid) + sizeof(uint16_t))) {
        int len;
        uuid *guid;

        guid = (uuid *)(ptr - sizeof(uuid));
        len = *(uint16_t *)(ptr - sizeof(uuid) - sizeof(uint16_t));

        /*
         * just in case the table is corrupt, wouldn't want to spin in
         * the zero case
         */
        if (len < (int)(sizeof(uuid) + sizeof(uint16_t))) {
            return false;
        } else if (len > tot_len) {
            return false;
        }



        /* @ptr points to the beginning of this entry */
        ptr -= len;
        if (qemu_uuid_is_equal(guid, &target_guid)) {
            if (data) {
                *data = ptr;
            }
            if (data_len) {
                *data_len = len - (int)(sizeof(uuid) + sizeof(uint16_t));
            }
            return true;
        }
        tot_len -= len;
    }
    return false;
}

/**
 * @brief find the reset vector GPA.
 *
 * @param [in] ovmf ovmf The pointer to memory corresponding to the OVMF flash file
 * @param [in] ovmf_size The size of the OVMF flash file
 * @param [in/out] addr The reset vector addr
 *
 * Return: 0 if found otherwise -1
 */
int ovmf_find_reset_vector(uint8_t *ovmf, size_t ovmf_size, uint32_t *addr)
{
    reset_info_block *info = NULL;
    uint8_t *table_data    = NULL;

    if (!ovmf || !addr) {
       printf("%s invalid ovmf or addr\n", __FUNCTION__);
       return -1;
    }

    /*
     * Initialize the address to zero. An address of zero with a successful
     * return code indicates that CSV2/CSV3 is not active.
     */
    *addr = 0;

    /*
     * Extract the AP reset vector for CSV2/3 guests by locating the Reset Vector GUID.
     */
    if (ovmf_table_find(ovmf, ovmf_size, RESET_BLOCK_GUID, &table_data, NULL)) {
        info = (reset_info_block*)table_data;
        if (info->reset_addr) {
            *addr = info->reset_addr;
            return 0;
        }
    }

    return -1;
}

/**
 * @brief find the hash table structure in GUID table.
 *
 * @param [in] ovmf ovmf The pointer to memory corresponding to the OVMF flash file
 * @param [in] ovmf_size The size of the OVMF flash file
 * @param [in/out] addr The hash table GPA
 *
 * Return: 0 if found otherwise -1
 */
int ovmf_find_hash_table_desc(uint8_t *ovmf, uint32_t ovmf_size, uint32_t *addr)
{
    hash_table_descriptor  *table = NULL;
    uint8_t *data                 = NULL;

    if (!ovmf || !addr) {
       printf("%s invalid ovmf or addr\n", __FUNCTION__);
       return -1;
    }

    if (ovmf_table_find(ovmf, ovmf_size, HASH_TABLE_GUID, &data, NULL)) {
        table = (hash_table_descriptor*)data;
        *addr = table->base;
        return 0;
    }

    printf("Error: OVMF has no hashes table GUID\n");
    return -1;
}

int fill_csv_hash_table(hash_table *ht, char *kernel_file, char *cmdline, char *initrd_file)
{
    uint8_t cmdline_hash[HASH_SIZE] = {0};
    uint8_t initrd_hash[HASH_SIZE] = {0};
    uint8_t kernel_hash[HASH_SIZE] = {0};
    size_t  cmdline_size;
    size_t  initrd_size;
    size_t  kernel_size;
    uint8_t *cmdline_buffer = NULL;
    uint8_t *initrd_buffer  = NULL;
    uint8_t *kernel_buffer  = NULL;
    uint8_t *p              = NULL;
    int ret = -1;
    uuid   tmp_guid;
    bool  cmdfile_exist = false;

    if (cmdline) {
        cmdfile_exist = file_exist(cmdline);
        if (cmdfile_exist)
            cmdline_size = get_file_length(cmdline);
    }
    if (cmdline == NULL || (cmdfile_exist && cmdline_size == 0)) {
        /* no cmdline passed in, must be '\0' */
        cmdline_buffer = (uint8_t*)malloc(1);
        cmdline_size = 1;
        cmdline_buffer[0] = '\0';
    } else {
        if (cmdfile_exist) {
            cmdline_buffer = (uint8_t*)malloc(cmdline_size);
            if (!cmdline_buffer)
                goto finish;
            read_file(cmdline, cmdline_buffer, cmdline_size);
            /*
             * if file end with 0xa(sometimes automatically added LF),
             * it should be replace by '\0'
             * otherwise add a extra '\0'
             */
            if (cmdline_buffer[cmdline_size-1] == 0xa)
                cmdline_buffer[cmdline_size-1] = 0x0;
            else {
                p = realloc(cmdline_buffer, cmdline_size + 1);
                cmdline_buffer = p;
                cmdline_size += 1;
                cmdline_buffer[cmdline_size-1] = 0x0;
            }
        } else {
            cmdline_buffer = (uint8_t*)cmdline;
            /* including the terminating-null */
            cmdline_size = strlen(cmdline) + 1;
        }
    }

    ret = sha256_hash(cmdline_buffer, cmdline_size, cmdline_hash);
    if (ret) {
        printf("cmdline hash error\n");
        goto finish;
    }

    /*
     * Calculate hash of initrd. If the user doesn't supply an initrd via
     * -initrd, an empty buffer will be used (initrd_size == 0).
     */
    initrd_size = get_file_length(initrd_file);
    if (initrd_size == 0) {
        printf("%s file doesn't exist or size is 0\n", initrd_file);
    } else {
        initrd_buffer = (uint8_t*)malloc(initrd_size);
        if (!initrd_buffer)
            goto finish;

        read_file(initrd_file, initrd_buffer, initrd_size);
    }

    ret = sha256_hash(initrd_buffer, initrd_size, initrd_hash);
    if (ret) {
        printf("initrd hash error\n");
        goto finish;
    }

    /* Calculate hash of the kernel */
    kernel_size = get_file_length(kernel_file);
    if (kernel_size == 0) {
        printf("%s file doesn't exist or size is 0,please check!\n", kernel_file);
        ret = -1;
        goto finish;
    } else {
        kernel_buffer = (uint8_t*)malloc(kernel_size);
        if (!kernel_buffer)
            goto finish;

        read_file(kernel_file, kernel_buffer, kernel_size);
    }

    ret = sha256_hash(kernel_buffer, kernel_size, kernel_hash);
    if (ret) {
        printf("kernel hash error\n");
        goto finish;
    }

    qemu_uuid_parse(HASH_TABLE_HEADER_GUID, &tmp_guid);
    ht->guid = qemu_uuid_bswap(tmp_guid);
    ht->len = sizeof(*ht);


    qemu_uuid_parse(CMDLINE_ENTRY_GUID, &tmp_guid);
    ht->cmdline.guid = qemu_uuid_bswap(tmp_guid);
    ht->cmdline.len = sizeof(ht->cmdline);
    memcpy(ht->cmdline.hash, cmdline_hash, sizeof(ht->cmdline.hash));

    qemu_uuid_parse(INITRD_ENTRY_GUID, &tmp_guid);
    ht->initrd.guid = qemu_uuid_bswap(tmp_guid);
    ht->initrd.len = sizeof(ht->initrd);
    memcpy(ht->initrd.hash, initrd_hash, sizeof(ht->initrd.hash));

    qemu_uuid_parse(KERNEL_ENTRY_GUID, &tmp_guid);
    ht->kernel.guid = qemu_uuid_bswap(tmp_guid);
    ht->kernel.len = sizeof(ht->kernel);
    memcpy(ht->kernel.hash, kernel_hash, sizeof(ht->kernel.hash));

    ret = 0;

finish:
    if (kernel_buffer)
        free(kernel_buffer);

    if (initrd_buffer)
        free(initrd_buffer);

    if (cmdfile_exist && cmdline_buffer)
        free(cmdline_buffer);

    return ret;
}
